// Cz Cranelift 后端 - FFI 接口
use cranelift::prelude::*;
use cranelift_module::{DataDescription, FuncId, Linkage, Module};
use cranelift_object::{ObjectBuilder, ObjectModule};
use std::collections::HashMap;

pub struct CraneliftBackend {
    module: ObjectModule,
    ctx: codegen::Context,
    functions: HashMap<String, FuncId>,
}

impl CraneliftBackend {
    pub fn new() -> Result<Self, String> {
        let isa = cranelift_native::builder()
            .map_err(|e| format!("Failed to create ISA builder: {}", e))?
            .finish(settings::Flags::new(settings::builder()))
            .map_err(|e| format!("Failed to create ISA: {}", e))?;

        let builder = ObjectBuilder::new(
            isa,
            "output",
            cranelift_module::default_libcall_names(),
        )
        .map_err(|e| format!("Failed to create object builder: {}", e))?;

        let module = ObjectModule::new(builder);
        let ctx = module.make_context();

        Ok(CraneliftBackend {
            module,
            ctx,
            functions: HashMap::new(),
        })
    }

    pub fn compile_function(
        &mut self,
        name: &str,
        params: &[(String, Type)],
        return_type: Option<Type>,
    ) -> Result<(), String> {
        // 创建函数签名
        let mut sig = self.module.make_signature();
        for (_, ty) in params {
            sig.params.push(AbiParam::new(*ty));
        }
        if let Some(ret_ty) = return_type {
            sig.returns.push(AbiParam::new(ret_ty));
        }

        // 声明函数
        let func_id = self
            .module
            .declare_function(name, Linkage::Export, &sig)
            .map_err(|e| format!("Failed to declare function: {}", e))?;

        self.functions.insert(name.to_string(), func_id);

        // 创建函数体
        self.ctx.func.signature = sig;
        self.ctx.func.name = cranelift_codegen::ir::ExternalName::user(0, func_id.as_u32());

        // 创建入口块
        let block = self.ctx.func.dfg.make_block();
        let mut builder = FunctionBuilder::new(&mut self.ctx.func, &mut self.ctx.func_ctx);
        builder.append_block_params_for_function_params(block);
        builder.switch_to_block(block);
        builder.seal_block(block);

        // 简单示例：返回第一个参数或 0
        if !params.is_empty() {
            let param_val = builder.block_params(block)[0];
            builder.ins().return_(&[param_val]);
        } else if return_type.is_some() {
            let zero = builder.ins().iconst(types::I32, 0);
            builder.ins().return_(&[zero]);
        } else {
            builder.ins().return_(&[]);
        }

        builder.finalize();

        // 定义函数
        self.module
            .define_function(func_id, &mut self.ctx)
            .map_err(|e| format!("Failed to define function: {}", e))?;

        self.module.clear_context(&mut self.ctx);

        Ok(())
    }

    pub fn emit_object(&mut self) -> Result<Vec<u8>, String> {
        let product = self
            .module
            .finish()
            .object
            .write()
            .map_err(|e| format!("Failed to write object: {}", e))?;

        Ok(product)
    }
}

// C FFI 接口
#[repr(C)]
pub struct CBackend {
    backend: *mut CraneliftBackend,
}

#[no_mangle]
pub extern "C" fn cranelift_backend_new() -> *mut CBackend {
    match CraneliftBackend::new() {
        Ok(backend) => {
            let c_backend = Box::new(CBackend {
                backend: Box::into_raw(Box::new(backend)),
            });
            Box::into_raw(c_backend)
        }
        Err(_) => std::ptr::null_mut(),
    }
}

#[no_mangle]
pub extern "C" fn cranelift_backend_free(backend: *mut CBackend) {
    if !backend.is_null() {
        unsafe {
            let c_backend = Box::from_raw(backend);
            if !c_backend.backend.is_null() {
                let _ = Box::from_raw(c_backend.backend);
            }
        }
    }
}

#[no_mangle]
pub extern "C" fn cranelift_compile_simple_function(
    backend: *mut CBackend,
    name: *const u8,
    name_len: usize,
) -> bool {
    if backend.is_null() || name.is_null() {
        return false;
    }

    unsafe {
        let c_backend = &mut *backend;
        let backend_ref = &mut *c_backend.backend;
        let name_slice = std::slice::from_raw_parts(name, name_len);
        let name_str = match std::str::from_utf8(name_slice) {
            Ok(s) => s,
            Err(_) => return false,
        };

        // 编译一个简单的函数：i32 -> i32
        backend_ref
            .compile_function(name_str, &[], Some(types::I32))
            .is_ok()
    }
}

#[no_mangle]
pub extern "C" fn cranelift_emit_object(
    backend: *mut CBackend,
    output: *mut *mut u8,
    output_len: *mut usize,
) -> bool {
    if backend.is_null() || output.is_null() || output_len.is_null() {
        return false;
    }

    unsafe {
        let c_backend = &mut *backend;
        let backend_ref = &mut *c_backend.backend;

        match backend_ref.emit_object() {
            Ok(data) => {
                let len = data.len();
                let ptr = Box::into_raw(data.into_boxed_slice()) as *mut u8;
                *output = ptr;
                *output_len = len;
                true
            }
            Err(_) => false,
        }
    }
}

#[no_mangle]
pub extern "C" fn cranelift_free_buffer(ptr: *mut u8, len: usize) {
    if !ptr.is_null() && len > 0 {
        unsafe {
            let _ = Vec::from_raw_parts(ptr, len, len);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_backend_creation() {
        let backend = CraneliftBackend::new();
        assert!(backend.is_ok());
    }

    #[test]
    fn test_simple_function() {
        let mut backend = CraneliftBackend::new().unwrap();
        let result = backend.compile_function("test_func", &[], Some(types::I32));
        assert!(result.is_ok());
    }
}
